import copy
import os.path as osp
import cv2
import numpy as np
import onnxruntime as ort
from glob import glob
from tqdm import tqdm
import os
import cv2
import base64
import json

def bbox2json(bboxes,categorys,imgs,image_path,json_dir):
    assert len(bboxes)==len(categorys)
    assert len(image_path)==len(bboxes)

    for i in range(len(image_path)):
        res={"version": "4.2.10",
                "flags": {},}
        json_path = os.path.join(json_dir, image_path[i].split('/')[-1].split('.png')[0] + '.json')
        with open(json_path,encoding='utf8',mode='w') as f:
            res['imagePath']=os.path.basename( image_path[i])
            bbox=bboxes[i]
            category=categorys[i]
            img=imgs[i]
            h,w,c=img.shape
            imag = cv2.imencode('.jpg', img)[1]
            base64_data = str(base64.b64encode(imag))[2:-1]
            #b64=base64.b64encode(img).decode('utf-8')
            res['imageData']=base64_data
            cur_res=[]
            #print(category)
            for j in range(len(bbox)):
                cur_poly=bbox[j]
                cur_res.append({'label':category[j],'points':[[poly[0],poly[1]] for poly in cur_poly],"group_id": 'null',"shape_type": "polygon","flags": {}})
            res['shapes']=cur_res
            # res['time_labeled']=13141516
            # res['labeled']='true'
            im=cv2.imread(image_path[i])
            h,w,c=im.shape
            res['imageHeight']=h
            res['imageWidth']=w
            res=json.dumps(res)
            f.write(res)
def read_json(json_file):
    data = open(json_file, "r")
    data = json.load(data)
    clses=[]
    polys=[]
    for info in data:
        name=info['name']
        idx=info['labelIdx']
        points=info['points']
        if name=="":
            clses.append(str(idx))
        else:
            clses.append(name)
        polys.append(points)
    return clses,polys

if __name__=='__main__':
    src_dir='/media/wsl/a9f0161f-7971-c843-8c81-c68049a0235a/DataSet/芜湖海螺/无人机/20220718/vis'
    eiseg_path=osp.join(src_dir,'label')
    files=glob(eiseg_path+'/*.json')
    for file in tqdm(files):
        clses,polys=read_json(file)
        im_path=osp.join(src_dir,os.path.basename(file).replace('json','png'))
        im=cv2.imread(im_path)
        bbox2json([polys], [clses], [im], [im_path], src_dir)
